import os
import mycv
import RotateImg
import utils
import random
import numpy as np
from tqdm import tqdm

path_Limit = r'/media/fang/TOSHIBA EXT/Temp/Circle/Small'
path_RemoveLimit = r'/media/fang/TOSHIBA EXT/Temp/Circle/RemoveLimit/SmallLabel'
path_UnLimit = r'/media/fang/TOSHIBA EXT/Temp/Circle/RemoveLimit/UnLimit'
path_out = r'/media/fang/Samsung USB/Small_from_Src'
max_shape_list = [[100, 100, 3], [150, 150, 3]]
out_img_shape = (300, 600, 3)
nums_out_img = 500
padding_min = 0.2
padding_max = 0.4

names_Limit, names_RemoveLimit, names_UnLimit = [], [], []
for root, dirs, files in os.walk(path_Limit):
    for file in files:
        if file.split('.')[-1] == 'jpg':
            names_Limit.append(root + '/' + file)
for root, dirs, files in os.walk(path_RemoveLimit):
    for file in files:
        if file.split('.')[-1] == 'jpg':
            names_RemoveLimit.append(root + '/' + file)
for root, dirs, files in os.walk(path_UnLimit):
    for file in files:
        if file.split('.')[-1] == 'jpg':
            names_UnLimit.append(root + '/' + file)
names_all = [names_Limit, names_RemoveLimit, names_UnLimit]

for idx_max_shape, max_shape in enumerate(max_shape_list):
    ContainsNum = int(out_img_shape[0] * out_img_shape[1] / max_shape[0] / max_shape[1])
    NameNum_offset = idx_max_shape * nums_out_img
    for NameNum in tqdm(range(nums_out_img)):
        Infos_output = [path_out + '/B%05d.xml' % (NameNum + NameNum_offset)]
        img_output = np.ones(out_img_shape, dtype=np.uint8) * 127
        for index_in_contain in range(ContainsNum):
            chosen_jpg = random.choice(random.choice(names_all))
            img = utils.imread(chosen_jpg)
            Infos = utils.ReadXml(chosen_jpg.replace('.jpg', '.xml'))

            try:
                cut_info_src = random.choice(Infos)
            except:
                print("Infos is empty sequence")
                print("file is" + chosen_jpg)
                print(Infos_output[0])
                continue
            padding_up = random.uniform(0, padding_max)
            padding_down = random.uniform(0, padding_max)
            padding_left = random.uniform(0, padding_max)
            padding_right = random.uniform(0, padding_max)

            cut_x = int(cut_info_src[0] - (cut_info_src[2] - cut_info_src[0]) * padding_left)
            cut_y = int(cut_info_src[1] - (cut_info_src[3] - cut_info_src[1]) * padding_up)
            cut_w = int((cut_info_src[2] - cut_info_src[0]) * (1 + padding_left + padding_right))
            cut_h = int((cut_info_src[3] - cut_info_src[1]) * (1 + padding_up + padding_down))

            if cut_x < 0:
                cut_x = 0
            if cut_y < 0:
                cut_y = 0
            if cut_x + cut_w > img.shape[1]:
                cut_x = img.shape[1] - cut_w
            if cut_y + cut_h > img.shape[0]:
                cut_y = img.shape[0] - cut_h

            x_start = index_in_contain % int(out_img_shape[1] / max_shape[1]) * max_shape[1]
            x_end = x_start + max_shape[1]
            y_start = int(index_in_contain / int(out_img_shape[1] / max_shape[1])) * max_shape[0]
            y_end = y_start + max_shape[0]

            cut_img = img[cut_y:cut_y + cut_h, cut_x:cut_x + cut_w]
            cut_info_dst = [[cut_info_src[0] - cut_x, cut_info_src[1] - cut_y, cut_info_src[2] - cut_x, cut_info_src[3] - cut_y, cut_info_src[4]]]

            rotate_img, rotate_info = RotateImg.RotateImgXml(cut_img, cut_info_dst, [-5, 5])

            resize_img, resize_info = utils.ResizeJpgAndXml(rotate_img, rotate_info, max_shape[1], max_shape[0])

            blur_img = mycv.motion_blur(resize_img, random.randint(1, 5))

            img_output[y_start:y_end, x_start:x_end] = blur_img
            Infos_output.append([resize_info[0][0] + x_start, resize_info[0][1] + y_start, resize_info[0][2] + x_start, resize_info[0][3] + y_start, resize_info[0][4]])

        utils.WriteXml(Infos_output, out_img_shape[1], out_img_shape[0])
        utils.imwrite(Infos_output[0].replace('xml', 'jpg'), img_output)
        pass
